# time_stepping.jl
# Advanced time stepping schemes for NSEMSolver with adaptive capabilities

"""
Advanced Time Stepping Methods for Navier-Stokes Equations

This module implements sophisticated time stepping schemes including:
- Embedded Runge-Kutta methods with error estimation
- Adaptive step size control with PI controllers  
- Multi-step methods (Adams-Bashforth, Adams-Moulton)
- IMEX (Implicit-Explicit) schemes for stiff problems
- Stability monitoring and automatic step rejection
"""

using LinearAlgebra
using Printf


"""
    AdaptiveTimestepper

Advanced adaptive time stepping with embedded error control.

# Fields
- `method::Symbol`: Time integration method (`:rk45`, `:rk23`, `:dopri5`, `:rkf45`)
- `tolerance::Float64`: Local error tolerance
- `safety_factor::Float64`: Safety factor for step size updates
- `min_dt::Float64`: Minimum allowed time step
- `max_dt::Float64`: Maximum allowed time step
- `max_step_ratio::Float64`: Maximum step size increase ratio
- `min_step_ratio::Float64`: Minimum step size decrease ratio
- `pi_controller::Bool`: Use PI controller for step size evolution
- `beta1::Float64`: Proportional gain for PI controller
- `beta2::Float64`: Integral gain for PI controller
- `facold::Float64`: Previous step size factor (PI controller state)
- `reject_threshold::Float64`: Threshold for step rejection
- `accept_threshold::Float64`: Threshold for step acceptance
- `verbose::Bool`: Enable verbose output
"""
Base.@kwdef mutable struct AdaptiveTimestepper
    method::Symbol = :rk45
    tolerance::Float64 = 1e-3
    safety_factor::Float64 = 0.9
    min_dt::Float64 = 1e-10
    max_dt::Float64 = 1.0
    max_step_ratio::Float64 = 5.0
    min_step_ratio::Float64 = 0.2
    pi_controller::Bool = true
    beta1::Float64 = 0.7  # Proportional gain
    beta2::Float64 = 0.4  # Integral gain  
    facold::Float64 = 1.0  # Previous factor
    reject_threshold::Float64 = 1.0
    accept_threshold::Float64 = 0.1
    verbose::Bool = false
end

"""
    TimeStepResult

Result from adaptive time step.
"""
struct TimeStepResult
    dt_new::Float64
    error_estimate::Float64
    accepted::Bool
    step_ratio::Float64
    iterations::Int
end

"""
    ButcherTableau

Butcher tableau for Runge-Kutta methods.
"""
struct ButcherTableau
    A::Matrix{Float64}  # RK matrix
    b::Vector{Float64}  # weights
    b_hat::Vector{Float64}  # embedded weights (for error estimation)
    c::Vector{Float64}  # nodes
    order::Int  # method order
    error_order::Int  # embedded method order
    name::String
end

"""
    adaptive_time_stepping(u::Array, v::Array, options::NSOptions) -> Float64

Compute adaptive time step based on CFL condition and solution characteristics.
(Legacy interface - use compute_stable_timestep for new code)
"""
function adaptive_time_stepping(u::Array, v::Array, options::NSOptions)
    # Compute maximum velocity magnitude
    u_max = maximum(abs, u)
    v_max = maximum(abs, v)
    vel_max = sqrt(u_max^2 + v_max^2)
    
    if vel_max > 0
        # CFL-based time step
        dx_min = 2.0 / (options.n_block * options.N)  # Minimum grid spacing
        dt_cfl = options.cfl * dx_min / vel_max
    else
        # Default time step if velocity is zero
        dt_cfl = 0.01
    end
    
    # Viscous time step constraint
    dx_min = 2.0 / (options.n_block * options.N)
    dt_visc = 0.5 * dx_min^2 / options.nu
    
    # Take the minimum of CFL and viscous constraints
    return min(dt_cfl, dt_visc)
end

"""
    runge_kutta_step!(u_new::Array, v_new::Array, u::Array, v::Array, 
                     dt::Float64, options::NSOptions)

Perform one Runge-Kutta time step for the Navier-Stokes equations.
"""
function runge_kutta_step!(u_new::Array, v_new::Array, u::Array, v::Array, 
                          dt::Float64, options::NSOptions)
    # Use default RK4 method
    tableau = get_butcher_tableau(:rk4)
    
    # Compute RK step (placeholder - would use actual NS operators)
    # This is a simplified implementation
    dudt = -options.nu * u  # Simple damping term
    dvdt = -options.nu * v
    
    # RK4 implementation
    k1_u = dt * dudt
    k1_v = dt * dvdt
    
    k2_u = dt * (-options.nu * (u + 0.5 * k1_u))
    k2_v = dt * (-options.nu * (v + 0.5 * k1_v))
    
    k3_u = dt * (-options.nu * (u + 0.5 * k2_u))
    k3_v = dt * (-options.nu * (v + 0.5 * k3_v))
    
    k4_u = dt * (-options.nu * (u + k3_u))
    k4_v = dt * (-options.nu * (v + k3_v))
    
    u_new[:] = u + (k1_u + 2*k2_u + 2*k3_u + k4_u) / 6
    v_new[:] = v + (k1_v + 2*k2_v + 2*k3_v + k4_v) / 6
    
    return nothing
end

"""
    adaptive_timestep!(u_new::Array, v_new::Array, u::Array, v::Array, 
                      dt::Float64, stepper::AdaptiveTimestepper, rhs_function, args...) -> TimeStepResult

Perform adaptive time step with embedded error estimation.
"""
function adaptive_timestep!(u_new::Array, v_new::Array, u::Array, v::Array, 
                           dt::Float64, stepper::AdaptiveTimestepper, rhs_function, args...)
    
    # Get Butcher tableau for chosen method
    tableau = get_butcher_tableau(stepper.method)
    
    # Compute embedded Runge-Kutta step
    u_high, v_high, u_low, v_low = embedded_rk_step(u, v, dt, tableau, rhs_function, args...)
    
    # Estimate local truncation error
    error_est = estimate_error(u_high, u_low, v_high, v_low, stepper.tolerance, stepper.tolerance)
    
    # Check if step should be accepted
    accepted = error_est <= stepper.reject_threshold
    
    if accepted
        # Accept step - use higher order solution
        u_new .= u_high
        v_new .= v_high
    end
    
    # Compute new step size
    dt_new, step_ratio = update_timestep(dt, error_est, stepper)
    
    if stepper.verbose && !accepted
        println("   ⚠️  Step rejected: dt=$(@sprintf("%.2e", dt)), error=$(@sprintf("%.2e", error_est))")
    end
    
    return TimeStepResult(dt_new, error_est, accepted, step_ratio, 1)
end

"""
    embedded_rk_step(u, v, dt, tableau, rhs_function, args...)

Compute embedded Runge-Kutta step for error estimation.
"""
function embedded_rk_step(u::Array, v::Array, dt::Float64, tableau::ButcherTableau, 
                         rhs_function, args...)
    
    s = length(tableau.c)  # Number of stages
    n = length(u)
    
    # Stage derivatives
    k_u = [zeros(size(u)) for _ in 1:s]
    k_v = [zeros(size(v)) for _ in 1:s]
    
    # Temporary arrays
    u_temp = similar(u)
    v_temp = similar(v)
    dudt = similar(u)
    dvdt = similar(v)
    
    # Compute stages
    for i in 1:s
        # Compute stage solution
        u_temp .= u
        v_temp .= v
        
        for j in 1:(i-1)
            u_temp .+= dt * tableau.A[i,j] * k_u[j]
            v_temp .+= dt * tableau.A[i,j] * k_v[j]
        end
        
        # Evaluate right-hand side
        rhs_function(dudt, dvdt, u_temp, v_temp, args...)
        
        k_u[i] .= dudt
        k_v[i] .= dvdt
    end
    
    # Higher order solution
    u_high = copy(u)
    v_high = copy(v)
    for i in 1:s
        u_high .+= dt * tableau.b[i] * k_u[i]
        v_high .+= dt * tableau.b[i] * k_v[i]
    end
    
    # Lower order solution (for error estimation)
    u_low = copy(u)
    v_low = copy(v)
    for i in 1:s
        u_low .+= dt * tableau.b_hat[i] * k_u[i]
        v_low .+= dt * tableau.b_hat[i] * k_v[i]
    end
    
    return u_high, v_high, u_low, v_low
end

"""
    estimate_error(u_high, u_low, v_high, v_low, atol, rtol)

Estimate local truncation error using embedded methods.
"""
function estimate_error(u_high::Array, u_low::Array, v_high::Array, v_low::Array, 
                       atol::Float64, rtol::Float64)
    
    error_u = u_high .- u_low
    error_v = v_high .- v_low
    
    # Compute scaled error norm
    n = length(u_high)
    scale_u = atol .+ rtol .* max.(abs.(u_high), abs.(u_low))
    scale_v = atol .+ rtol .* max.(abs.(v_high), abs.(v_low))
    
    # Root mean square error
    error_norm = sqrt((sum((error_u ./ scale_u).^2) + sum((error_v ./ scale_v).^2)) / (2*n))
    
    return error_norm
end

"""
    update_timestep(dt_old::Float64, error_ratio::Float64, stepper::AdaptiveTimestepper)

Update time step size based on error estimate and controller strategy.
"""
function update_timestep(dt_old::Float64, error_ratio::Float64, stepper::AdaptiveTimestepper)
    
    if error_ratio <= 0
        # Very small error - increase step size significantly
        factor = stepper.max_step_ratio
    else
        # Standard PI controller or simple strategy
        if stepper.pi_controller
            # PI controller for smooth step size evolution
            factor = stepper.safety_factor * (1/error_ratio)^stepper.beta1 * stepper.facold^stepper.beta2
        else
            # Simple controller
            factor = stepper.safety_factor * (1/error_ratio)^(1/(stepper.method == :rk23 ? 3 : 5))
        end
        
        # Limit step size changes
        factor = clamp(factor, stepper.min_step_ratio, stepper.max_step_ratio)
    end
    
    # Compute new step size
    dt_new = clamp(dt_old * factor, stepper.min_dt, stepper.max_dt)
    
    # Update PI controller state
    if stepper.pi_controller
        stepper.facold = factor
    end
    
    return dt_new, factor
end

"""
    get_butcher_tableau(method::Symbol) -> ButcherTableau

Get Butcher tableau for specified RK method.
"""
function get_butcher_tableau(method::Symbol)
    
    if method == :rk23
        # Bogacki-Shampine 2(3) method
        A = [0.0 0.0 0.0 0.0;
             0.5 0.0 0.0 0.0;
             0.0 0.75 0.0 0.0;
             2.0/9.0 1.0/3.0 4.0/9.0 0.0]
        b = [2.0/9.0, 1.0/3.0, 4.0/9.0, 0.0]
        b_hat = [7.0/24.0, 1.0/4.0, 1.0/3.0, 1.0/8.0]  # 3rd order
        c = [0.0, 0.5, 0.75, 1.0]
        return ButcherTableau(A, b, b_hat, c, 2, 3, "Bogacki-Shampine 2(3)")
        
    elseif method == :rk45 || method == :rkf45
        # Runge-Kutta-Fehlberg 4(5) method
        A = [0.0 0.0 0.0 0.0 0.0 0.0;
             1.0/4.0 0.0 0.0 0.0 0.0 0.0;
             3.0/32.0 9.0/32.0 0.0 0.0 0.0 0.0;
             1932.0/2197.0 -7200.0/2197.0 7296.0/2197.0 0.0 0.0 0.0;
             439.0/216.0 -8.0 3680.0/513.0 -845.0/4104.0 0.0 0.0;
             -8.0/27.0 2.0 -3544.0/2565.0 1859.0/4104.0 -11.0/40.0 0.0]
        b = [16.0/135.0, 0.0, 6656.0/12825.0, 28561.0/56430.0, -9.0/50.0, 2.0/55.0]
        b_hat = [25.0/216.0, 0.0, 1408.0/2565.0, 2197.0/4104.0, -1.0/5.0, 0.0]  # 4th order
        c = [0.0, 1.0/4.0, 3.0/8.0, 12.0/13.0, 1.0, 1.0/2.0]
        return ButcherTableau(A, b, b_hat, c, 5, 4, "Runge-Kutta-Fehlberg 4(5)")
        
    elseif method == :dopri5
        # Dormand-Prince 5(4) method
        A = [0.0 0.0 0.0 0.0 0.0 0.0 0.0;
             1.0/5.0 0.0 0.0 0.0 0.0 0.0 0.0;
             3.0/40.0 9.0/40.0 0.0 0.0 0.0 0.0 0.0;
             44.0/45.0 -56.0/15.0 32.0/9.0 0.0 0.0 0.0 0.0;
             19372.0/6561.0 -25360.0/2187.0 64448.0/6561.0 -212.0/729.0 0.0 0.0 0.0;
             9017.0/3168.0 -355.0/33.0 46732.0/5247.0 49.0/176.0 -5103.0/18656.0 0.0 0.0;
             35.0/384.0 0.0 500.0/1113.0 125.0/192.0 -2187.0/6784.0 11.0/84.0 0.0]
        b = [35.0/384.0, 0.0, 500.0/1113.0, 125.0/192.0, -2187.0/6784.0, 11.0/84.0, 0.0]
        b_hat = [5179.0/57600.0, 0.0, 7571.0/16695.0, 393.0/640.0, -92097.0/339200.0, 187.0/2100.0, 1.0/40.0]
        c = [0.0, 1.0/5.0, 3.0/10.0, 4.0/5.0, 8.0/9.0, 1.0, 1.0]
        return ButcherTableau(A, b, b_hat, c, 5, 4, "Dormand-Prince 5(4)")
        
    elseif method == :rk4
        # Classical 4th order RK (no embedding)
        A = [0.0 0.0 0.0 0.0;
             0.5 0.0 0.0 0.0;
             0.0 0.5 0.0 0.0;
             0.0 0.0 1.0 0.0]
        b = [1.0/6.0, 1.0/3.0, 1.0/3.0, 1.0/6.0]
        b_hat = b  # Same as b (no error control)
        c = [0.0, 0.5, 0.5, 1.0]
        return ButcherTableau(A, b, b_hat, c, 4, 4, "Classical RK4")
        
    else
        error("Unknown RK method: $method")
    end
end

"""
    compute_stable_timestep(u::Array, v::Array, multidomain, options::NSOptions; 
                           cfl_safety::Float64=0.8) -> Float64

Compute stable time step considering CFL and diffusive constraints.
"""
function compute_stable_timestep(u::Array, v::Array, multidomain, options::NSOptions; 
                                cfl_safety::Float64=0.8)
    
    # Maximum velocity magnitude
    u_max = maximum(abs, u)
    v_max = maximum(abs, v)
    vel_max = sqrt(u_max^2 + v_max^2)
    
    # Characteristic mesh size
    h_min = 2.0 / (multidomain.n_block * multidomain.n)  # Minimum element size
    
    # CFL constraint for convection
    if vel_max > 1e-12
        dt_cfl = cfl_safety * h_min / vel_max
    else
        dt_cfl = Inf
    end
    
    # Diffusive constraint (more restrictive for high-order methods)
    dt_diff = cfl_safety * h_min^2 / (2 * multidomain.dim * options.nu)
    
    # Stability constraint for spectral elements
    dt_spec = cfl_safety * h_min^2 / (multidomain.n^2 * options.nu)
    
    # Take minimum of all constraints
    dt_stable = min(dt_cfl, dt_diff, dt_spec)
    
    # Ensure reasonable bounds
    dt_stable = max(dt_stable, 1e-12)
    dt_stable = min(dt_stable, 0.1)  # Maximum reasonable step
    
    return dt_stable
end

"""
    imex_step!(u_new::Array, v_new::Array, u::Array, v::Array, dt::Float64,
              explicit_rhs!, implicit_rhs!, linear_solver, options::NSOptions)

IMEX (Implicit-Explicit) time step for stiff systems.
"""
function imex_step!(u_new::Array, v_new::Array, u::Array, v::Array, dt::Float64,
                   explicit_rhs!, implicit_rhs!, linear_solver, options::NSOptions)
    
    # Explicit treatment of convection, implicit treatment of diffusion
    
    # Compute explicit (convection) terms
    explicit_u = similar(u)
    explicit_v = similar(v)
    explicit_rhs!(explicit_u, explicit_v, u, v)
    
    # Set up implicit system: (I - dt*J)*u_new = u + dt*explicit_u
    # where J is the Jacobian of implicit terms
    
    # Right-hand side
    rhs_u = u .+ dt .* explicit_u
    rhs_v = v .+ dt .* explicit_v
    
    # Solve implicit system (placeholder - would use actual linear solver)
    # (I - dt*viscous_operator) * u_new = rhs_u
    
    # For now, simple backward Euler approximation
    implicit_factor = 1.0 + dt * options.nu
    u_new .= rhs_u ./ implicit_factor
    v_new .= rhs_v ./ implicit_factor
    
    return nothing
end

"""
    adams_bashforth_step!(u_new::Array, v_new::Array, history::Vector, dt::Float64, order::Int)

Adams-Bashforth multi-step method.
"""
function adams_bashforth_step!(u_new::Array, v_new::Array, history::Vector, dt::Float64, order::Int)
    
    if length(history) < order
        error("Insufficient history for AB$order method")
    end
    
    # Adams-Bashforth coefficients
    if order == 1
        coeffs = [1.0]
    elseif order == 2
        coeffs = [3.0/2.0, -1.0/2.0]
    elseif order == 3
        coeffs = [23.0/12.0, -16.0/12.0, 5.0/12.0]
    elseif order == 4
        coeffs = [55.0/24.0, -59.0/24.0, 37.0/24.0, -9.0/24.0]
    else
        error("Adams-Bashforth order $order not implemented")
    end
    
    # Get current state
    u, v = history[end]
    u_new .= u
    v_new .= v
    
    # Add multi-step contributions
    for i in 1:order
        u_prev, v_prev, dudt_prev, dvdt_prev = history[end - i + 1]
        u_new .+= dt .* coeffs[i] .* dudt_prev
        v_new .+= dt .* coeffs[i] .* dvdt_prev
    end
    
    return nothing
end

"""
    create_adaptive_timestepper(method::Symbol; kwargs...) -> AdaptiveTimestepper

Create adaptive time stepper with specified method and options.
"""
function create_adaptive_timestepper(method::Symbol; kwargs...)
    return AdaptiveTimestepper(; method=method, kwargs...)
end

"""
    monitor_stability(u::Array, v::Array, dt::Float64, options::NSOptions) -> Bool

Monitor solution stability and detect blow-up.
"""
function monitor_stability(u::Array, v::Array, dt::Float64, options::NSOptions)
    
    # Check for NaN or Inf
    if any(isnan, u) || any(isnan, v) || any(isinf, u) || any(isinf, v)
        @warn "NaN or Inf detected in solution"
        return false
    end
    
    # Check for excessive growth
    u_max = maximum(abs, u)
    v_max = maximum(abs, v)
    
    if u_max > 1e6 || v_max > 1e6
        @warn "Solution appears to be growing without bound: u_max=$u_max, v_max=$v_max"
        return false
    end
    
    # Check CFL constraint
    vel_max = sqrt(u_max^2 + v_max^2)
    h_min = 2.0 / (options.n_block * options.N)
    cfl_actual = vel_max * dt / h_min
    
    if cfl_actual > 2.0  # Significantly above stability limit
        @warn "CFL number too large: CFL=$cfl_actual, dt=$dt"
        return false
    end
    
    return true
end